{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Fully Convolutional Networks for Change Detection\n",
    "\n",
    "Example code for training the network presented in the paper:\n",
    "\n",
    "```\n",
    "Daudt, R.C., Le Saux, B. and Boulch, A., 2018, October. Fully convolutional siamese networks for change detection. In 2018 25th IEEE International Conference on Image Processing (ICIP) (pp. 4063-4067). IEEE.\n",
    "```\n",
    "\n",
    "Code uses the OSCD dataset:\n",
    "\n",
    "```\n",
    "Daudt, R.C., Le Saux, B., Boulch, A. and Gousseau, Y., 2018, July. Urban change detection for multispectral earth observation using convolutional neural networks. In IGARSS 2018-2018 IEEE International Geoscience and Remote Sensing Symposium (pp. 2115-2118). IEEE.\n",
    "```\n",
    "\n",
    "\n",
    "FresUNet architecture from paper:\n",
    "\n",
    "```\n",
    "Daudt, R.C., Le Saux, B., Boulch, A. and Gousseau, Y., 2019. Multitask learning for large-scale semantic change detection. Computer Vision and Image Understanding, 187, p.102783.\n",
    "```\n",
    "\n",
    "Please consider all relevant papers if you use this code."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Rodrigo Daudt\n",
    "# rcdaudt.github.io\n",
    "# rodrigo.daudt@onera.fr"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Imports\n",
    "\n",
    "# PyTorch\n",
    "import torch\n",
    "import torch.nn as nn\n",
    "from torch.utils.data import Dataset, DataLoader\n",
    "from torch.autograd import Variable\n",
    "import torchvision.transforms as tr\n",
    "\n",
    "# Models\n",
    "from unet import Unet\n",
    "from siamunet_conc import SiamUnet_conc\n",
    "from siamunet_diff import SiamUnet_diff\n",
    "from fresunet import FresUNet\n",
    "\n",
    "# Other\n",
    "import os\n",
    "import numpy as np\n",
    "import random\n",
    "from skimage import io\n",
    "from scipy.ndimage import zoom\n",
    "import matplotlib.pyplot as plt\n",
    "%matplotlib inline\n",
    "from tqdm import tqdm as tqdm\n",
    "from pandas import read_csv\n",
    "from math import floor, ceil, sqrt, exp\n",
    "from IPython import display\n",
    "import time\n",
    "from itertools import chain\n",
    "import time\n",
    "import warnings\n",
    "from pprint import pprint\n",
    "\n",
    "\n",
    "\n",
    "print('IMPORTS OK')\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Global Variables' Definitions\n",
    "\n",
    "PATH_TO_DATASET = './OSCD/'\n",
    "IS_PROTOTYPE = False\n",
    "\n",
    "FP_MODIFIER = 10 # Tuning parameter, use 1 if unsure\n",
    "\n",
    "BATCH_SIZE = 32\n",
    "PATCH_SIDE = 96\n",
    "N_EPOCHS = 50\n",
    "\n",
    "NORMALISE_IMGS = True\n",
    "\n",
    "TRAIN_STRIDE = int(PATCH_SIDE/2) - 1\n",
    "\n",
    "TYPE = 3 # 0-RGB | 1-RGBIr | 2-All bands s.t. resulution <= 20m | 3-All bands\n",
    "\n",
    "LOAD_TRAINED = False\n",
    "\n",
    "DATA_AUG = True\n",
    "\n",
    "\n",
    "print('DEFINITIONS OK')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Functions\n",
    "\n",
    "def adjust_shape(I, s):\n",
    "    \"\"\"Adjust shape of grayscale image I to s.\"\"\"\n",
    "    \n",
    "    # crop if necesary\n",
    "    I = I[:s[0],:s[1]]\n",
    "    si = I.shape\n",
    "    \n",
    "    # pad if necessary \n",
    "    p0 = max(0,s[0] - si[0])\n",
    "    p1 = max(0,s[1] - si[1])\n",
    "    \n",
    "    return np.pad(I,((0,p0),(0,p1)),'edge')\n",
    "    \n",
    "\n",
    "def read_sentinel_img(path):\n",
    "    \"\"\"Read cropped Sentinel-2 image: RGB bands.\"\"\"\n",
    "    im_name = os.listdir(path)[0][:-7]\n",
    "    r = io.imread(path + im_name + \"B04.tif\")\n",
    "    g = io.imread(path + im_name + \"B03.tif\")\n",
    "    b = io.imread(path + im_name + \"B02.tif\")\n",
    "    \n",
    "    I = np.stack((r,g,b),axis=2).astype('float')\n",
    "    \n",
    "    if NORMALISE_IMGS:\n",
    "        I = (I - I.mean()) / I.std()\n",
    "\n",
    "    return I\n",
    "\n",
    "def read_sentinel_img_4(path):\n",
    "    \"\"\"Read cropped Sentinel-2 image: RGB and NIR bands.\"\"\"\n",
    "    im_name = os.listdir(path)[0][:-7]\n",
    "    r = io.imread(path + im_name + \"B04.tif\")\n",
    "    g = io.imread(path + im_name + \"B03.tif\")\n",
    "    b = io.imread(path + im_name + \"B02.tif\")\n",
    "    nir = io.imread(path + im_name + \"B08.tif\")\n",
    "    \n",
    "    I = np.stack((r,g,b,nir),axis=2).astype('float')\n",
    "    \n",
    "    if NORMALISE_IMGS:\n",
    "        I = (I - I.mean()) / I.std()\n",
    "\n",
    "    return I\n",
    "\n",
    "def read_sentinel_img_leq20(path):\n",
    "    \"\"\"Read cropped Sentinel-2 image: bands with resolution less than or equals to 20m.\"\"\"\n",
    "    im_name = os.listdir(path)[0][:-7]\n",
    "    \n",
    "    r = io.imread(path + im_name + \"B04.tif\")\n",
    "    s = r.shape\n",
    "    g = io.imread(path + im_name + \"B03.tif\")\n",
    "    b = io.imread(path + im_name + \"B02.tif\")\n",
    "    nir = io.imread(path + im_name + \"B08.tif\")\n",
    "    \n",
    "    ir1 = adjust_shape(zoom(io.imread(path + im_name + \"B05.tif\"),2),s)\n",
    "    ir2 = adjust_shape(zoom(io.imread(path + im_name + \"B06.tif\"),2),s)\n",
    "    ir3 = adjust_shape(zoom(io.imread(path + im_name + \"B07.tif\"),2),s)\n",
    "    nir2 = adjust_shape(zoom(io.imread(path + im_name + \"B8A.tif\"),2),s)\n",
    "    swir2 = adjust_shape(zoom(io.imread(path + im_name + \"B11.tif\"),2),s)\n",
    "    swir3 = adjust_shape(zoom(io.imread(path + im_name + \"B12.tif\"),2),s)\n",
    "    \n",
    "    I = np.stack((r,g,b,nir,ir1,ir2,ir3,nir2,swir2,swir3),axis=2).astype('float')\n",
    "    \n",
    "    if NORMALISE_IMGS:\n",
    "        I = (I - I.mean()) / I.std()\n",
    "\n",
    "    return I\n",
    "\n",
    "def read_sentinel_img_leq60(path):\n",
    "    \"\"\"Read cropped Sentinel-2 image: all bands.\"\"\"\n",
    "    im_name = os.listdir(path)[0][:-7]\n",
    "    \n",
    "    r = io.imread(path + im_name + \"B04.tif\")\n",
    "    s = r.shape\n",
    "    g = io.imread(path + im_name + \"B03.tif\")\n",
    "    b = io.imread(path + im_name + \"B02.tif\")\n",
    "    nir = io.imread(path + im_name + \"B08.tif\")\n",
    "    \n",
    "    ir1 = adjust_shape(zoom(io.imread(path + im_name + \"B05.tif\"),2),s)\n",
    "    ir2 = adjust_shape(zoom(io.imread(path + im_name + \"B06.tif\"),2),s)\n",
    "    ir3 = adjust_shape(zoom(io.imread(path + im_name + \"B07.tif\"),2),s)\n",
    "    nir2 = adjust_shape(zoom(io.imread(path + im_name + \"B8A.tif\"),2),s)\n",
    "    swir2 = adjust_shape(zoom(io.imread(path + im_name + \"B11.tif\"),2),s)\n",
    "    swir3 = adjust_shape(zoom(io.imread(path + im_name + \"B12.tif\"),2),s)\n",
    "    \n",
    "    uv = adjust_shape(zoom(io.imread(path + im_name + \"B01.tif\"),6),s)\n",
    "    wv = adjust_shape(zoom(io.imread(path + im_name + \"B09.tif\"),6),s)\n",
    "    swirc = adjust_shape(zoom(io.imread(path + im_name + \"B10.tif\"),6),s)\n",
    "    \n",
    "    I = np.stack((r,g,b,nir,ir1,ir2,ir3,nir2,swir2,swir3,uv,wv,swirc),axis=2).astype('float')\n",
    "    \n",
    "    if NORMALISE_IMGS:\n",
    "        I = (I - I.mean()) / I.std()\n",
    "\n",
    "    return I\n",
    "\n",
    "def read_sentinel_img_trio(path):\n",
    "    \"\"\"Read cropped Sentinel-2 image pair and change map.\"\"\"\n",
    "#     read images\n",
    "    if TYPE == 0:\n",
    "        I1 = read_sentinel_img(path + '/imgs_1/')\n",
    "        I2 = read_sentinel_img(path + '/imgs_2/')\n",
    "    elif TYPE == 1:\n",
    "        I1 = read_sentinel_img_4(path + '/imgs_1/')\n",
    "        I2 = read_sentinel_img_4(path + '/imgs_2/')\n",
    "    elif TYPE == 2:\n",
    "        I1 = read_sentinel_img_leq20(path + '/imgs_1/')\n",
    "        I2 = read_sentinel_img_leq20(path + '/imgs_2/')\n",
    "    elif TYPE == 3:\n",
    "        I1 = read_sentinel_img_leq60(path + '/imgs_1/')\n",
    "        I2 = read_sentinel_img_leq60(path + '/imgs_2/')\n",
    "        \n",
    "    cm = io.imread(path + '/cm/cm.png', as_gray=True) != 0\n",
    "    \n",
    "    # crop if necessary\n",
    "    s1 = I1.shape\n",
    "    s2 = I2.shape\n",
    "    I2 = np.pad(I2,((0, s1[0] - s2[0]), (0, s1[1] - s2[1]), (0,0)),'edge')\n",
    "    \n",
    "    \n",
    "    return I1, I2, cm\n",
    "\n",
    "\n",
    "\n",
    "def reshape_for_torch(I):\n",
    "    \"\"\"Transpose image for PyTorch coordinates.\"\"\"\n",
    "#     out = np.swapaxes(I,1,2)\n",
    "#     out = np.swapaxes(out,0,1)\n",
    "#     out = out[np.newaxis,:]\n",
    "    out = I.transpose((2, 0, 1))\n",
    "    return torch.from_numpy(out)\n",
    "\n",
    "\n",
    "\n",
    "class ChangeDetectionDataset(Dataset):\n",
    "    \"\"\"Change Detection dataset class, used for both training and test data.\"\"\"\n",
    "\n",
    "    def __init__(self, path, train = True, patch_side = 96, stride = None, use_all_bands = False, transform=None):\n",
    "        \"\"\"\n",
    "        Args:\n",
    "            csv_file (string): Path to the csv file with annotations.\n",
    "            root_dir (string): Directory with all the images.\n",
    "            transform (callable, optional): Optional transform to be applied\n",
    "                on a sample.\n",
    "        \"\"\"\n",
    "        \n",
    "        # basics\n",
    "        self.transform = transform\n",
    "        self.path = path\n",
    "        self.patch_side = patch_side\n",
    "        if not stride:\n",
    "            self.stride = 1\n",
    "        else:\n",
    "            self.stride = stride\n",
    "        \n",
    "        if train:\n",
    "            fname = 'train.txt'\n",
    "        else:\n",
    "            fname = 'test.txt'\n",
    "        \n",
    "#         print(path + fname)\n",
    "        self.names = read_csv(path + fname).columns\n",
    "        self.n_imgs = self.names.shape[0]\n",
    "        \n",
    "        n_pix = 0\n",
    "        true_pix = 0\n",
    "        \n",
    "        \n",
    "        # load images\n",
    "        self.imgs_1 = {}\n",
    "        self.imgs_2 = {}\n",
    "        self.change_maps = {}\n",
    "        self.n_patches_per_image = {}\n",
    "        self.n_patches = 0\n",
    "        self.patch_coords = []\n",
    "        for im_name in tqdm(self.names):\n",
    "            # load and store each image\n",
    "            I1, I2, cm = read_sentinel_img_trio(self.path + im_name)\n",
    "            self.imgs_1[im_name] = reshape_for_torch(I1)\n",
    "            self.imgs_2[im_name] = reshape_for_torch(I2)\n",
    "            self.change_maps[im_name] = cm\n",
    "            \n",
    "            s = cm.shape\n",
    "            n_pix += np.prod(s)\n",
    "            true_pix += cm.sum()\n",
    "            \n",
    "            # calculate the number of patches\n",
    "            s = self.imgs_1[im_name].shape\n",
    "            n1 = ceil((s[1] - self.patch_side + 1) / self.stride)\n",
    "            n2 = ceil((s[2] - self.patch_side + 1) / self.stride)\n",
    "            n_patches_i = n1 * n2\n",
    "            self.n_patches_per_image[im_name] = n_patches_i\n",
    "            self.n_patches += n_patches_i\n",
    "            \n",
    "            # generate path coordinates\n",
    "            for i in range(n1):\n",
    "                for j in range(n2):\n",
    "                    # coordinates in (x1, x2, y1, y2)\n",
    "                    current_patch_coords = (im_name, \n",
    "                                    [self.stride*i, self.stride*i + self.patch_side, self.stride*j, self.stride*j + self.patch_side],\n",
    "                                    [self.stride*(i + 1), self.stride*(j + 1)])\n",
    "                    self.patch_coords.append(current_patch_coords)\n",
    "                    \n",
    "        self.weights = [ FP_MODIFIER * 2 * true_pix / n_pix, 2 * (n_pix - true_pix) / n_pix]\n",
    "        \n",
    "        \n",
    "\n",
    "    def get_img(self, im_name):\n",
    "        return self.imgs_1[im_name], self.imgs_2[im_name], self.change_maps[im_name]\n",
    "\n",
    "    def __len__(self):\n",
    "        return self.n_patches\n",
    "\n",
    "    def __getitem__(self, idx):\n",
    "        current_patch_coords = self.patch_coords[idx]\n",
    "        im_name = current_patch_coords[0]\n",
    "        limits = current_patch_coords[1]\n",
    "        centre = current_patch_coords[2]\n",
    "        \n",
    "        I1 = self.imgs_1[im_name][:, limits[0]:limits[1], limits[2]:limits[3]]\n",
    "        I2 = self.imgs_2[im_name][:, limits[0]:limits[1], limits[2]:limits[3]]\n",
    "        \n",
    "        label = self.change_maps[im_name][limits[0]:limits[1], limits[2]:limits[3]]\n",
    "        label = torch.from_numpy(1*np.array(label)).float()\n",
    "        \n",
    "        sample = {'I1': I1, 'I2': I2, 'label': label}\n",
    "        \n",
    "        if self.transform:\n",
    "            sample = self.transform(sample)\n",
    "\n",
    "        return sample\n",
    "\n",
    "\n",
    "\n",
    "class RandomFlip(object):\n",
    "    \"\"\"Flip randomly the images in a sample.\"\"\"\n",
    "\n",
    "#     def __init__(self):\n",
    "#         return\n",
    "\n",
    "    def __call__(self, sample):\n",
    "        I1, I2, label = sample['I1'], sample['I2'], sample['label']\n",
    "        \n",
    "        if random.random() > 0.5:\n",
    "            I1 =  I1.numpy()[:,:,::-1].copy()\n",
    "            I1 = torch.from_numpy(I1)\n",
    "            I2 =  I2.numpy()[:,:,::-1].copy()\n",
    "            I2 = torch.from_numpy(I2)\n",
    "            label =  label.numpy()[:,::-1].copy()\n",
    "            label = torch.from_numpy(label)\n",
    "\n",
    "        return {'I1': I1, 'I2': I2, 'label': label}\n",
    "\n",
    "\n",
    "\n",
    "class RandomRot(object):\n",
    "    \"\"\"Rotate randomly the images in a sample.\"\"\"\n",
    "\n",
    "#     def __init__(self):\n",
    "#         return\n",
    "\n",
    "    def __call__(self, sample):\n",
    "        I1, I2, label = sample['I1'], sample['I2'], sample['label']\n",
    "        \n",
    "        n = random.randint(0, 3)\n",
    "        if n:\n",
    "            I1 =  sample['I1'].numpy()\n",
    "            I1 = np.rot90(I1, n, axes=(1, 2)).copy()\n",
    "            I1 = torch.from_numpy(I1)\n",
    "            I2 =  sample['I2'].numpy()\n",
    "            I2 = np.rot90(I2, n, axes=(1, 2)).copy()\n",
    "            I2 = torch.from_numpy(I2)\n",
    "            label =  sample['label'].numpy()\n",
    "            label = np.rot90(label, n, axes=(0, 1)).copy()\n",
    "            label = torch.from_numpy(label)\n",
    "\n",
    "        return {'I1': I1, 'I2': I2, 'label': label}\n",
    "\n",
    "\n",
    "\n",
    "\n",
    "\n",
    "print('UTILS OK')\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Dataset\n",
    "\n",
    "\n",
    "if DATA_AUG:\n",
    "    data_transform = tr.Compose([RandomFlip(), RandomRot()])\n",
    "else:\n",
    "    data_transform = None\n",
    "\n",
    "\n",
    "        \n",
    "\n",
    "train_dataset = ChangeDetectionDataset(PATH_TO_DATASET, train = True, stride = TRAIN_STRIDE, transform=data_transform)\n",
    "weights = torch.FloatTensor(train_dataset.weights).cuda()\n",
    "print(weights)\n",
    "train_loader = DataLoader(train_dataset, batch_size = BATCH_SIZE, shuffle = True, num_workers = 4)\n",
    "\n",
    "test_dataset = ChangeDetectionDataset(PATH_TO_DATASET, train = False, stride = TRAIN_STRIDE)\n",
    "test_loader = DataLoader(test_dataset, batch_size = BATCH_SIZE, shuffle = True, num_workers = 4)\n",
    "\n",
    "\n",
    "print('DATASETS OK')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# print(weights)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# 0-RGB | 1-RGBIr | 2-All bands s.t. resulution <= 20m | 3-All bands\n",
    "\n",
    "if TYPE == 0:\n",
    "#     net, net_name = Unet(2*3, 2), 'FC-EF'\n",
    "#     net, net_name = SiamUnet_conc(3, 2), 'FC-Siam-conc'\n",
    "#     net, net_name = SiamUnet_diff(3, 2), 'FC-Siam-diff'\n",
    "    net, net_name = FresUNet(2*3, 2), 'FresUNet'\n",
    "elif TYPE == 1:\n",
    "#     net, net_name = Unet(2*4, 2), 'FC-EF'\n",
    "#     net, net_name = SiamUnet_conc(4, 2), 'FC-Siam-conc'\n",
    "#     net, net_name = SiamUnet_diff(4, 2), 'FC-Siam-diff'\n",
    "    net, net_name = FresUNet(2*4, 2), 'FresUNet'\n",
    "elif TYPE == 2:\n",
    "#     net, net_name = Unet(2*10, 2), 'FC-EF'\n",
    "#     net, net_name = SiamUnet_conc(10, 2), 'FC-Siam-conc'\n",
    "#     net, net_name = SiamUnet_diff(10, 2), 'FC-Siam-diff'\n",
    "    net, net_name = FresUNet(2*10, 2), 'FresUNet'\n",
    "elif TYPE == 3:\n",
    "#     net, net_name = Unet(2*13, 2), 'FC-EF'\n",
    "#     net, net_name = SiamUnet_conc(13, 2), 'FC-Siam-conc'\n",
    "#     net, net_name = SiamUnet_diff(13, 2), 'FC-Siam-diff'\n",
    "    net, net_name = FresUNet(2*13, 2), 'FresUNet'\n",
    "\n",
    "\n",
    "net.cuda()\n",
    "\n",
    "criterion = nn.NLLLoss(weight=weights) # to be used with logsoftmax output\n",
    "\n",
    "print('NETWORK OK')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def count_parameters(model):\n",
    "    return sum(p.numel() for p in model.parameters() if p.requires_grad)\n",
    "\n",
    "print('Number of trainable parameters:', count_parameters(net))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# net.load_state_dict(torch.load('net-best_epoch-1_fm-0.7394933126157746.pth.tar'))\n",
    "\n",
    "def train(n_epochs = N_EPOCHS, save = True):\n",
    "    t = np.linspace(1, n_epochs, n_epochs)\n",
    "    \n",
    "    epoch_train_loss = 0 * t\n",
    "    epoch_train_accuracy = 0 * t\n",
    "    epoch_train_change_accuracy = 0 * t\n",
    "    epoch_train_nochange_accuracy = 0 * t\n",
    "    epoch_train_precision = 0 * t\n",
    "    epoch_train_recall = 0 * t\n",
    "    epoch_train_Fmeasure = 0 * t\n",
    "    epoch_test_loss = 0 * t\n",
    "    epoch_test_accuracy = 0 * t\n",
    "    epoch_test_change_accuracy = 0 * t\n",
    "    epoch_test_nochange_accuracy = 0 * t\n",
    "    epoch_test_precision = 0 * t\n",
    "    epoch_test_recall = 0 * t\n",
    "    epoch_test_Fmeasure = 0 * t\n",
    "    \n",
    "#     mean_acc = 0\n",
    "#     best_mean_acc = 0\n",
    "    fm = 0\n",
    "    best_fm = 0\n",
    "    \n",
    "    lss = 1000\n",
    "    best_lss = 1000\n",
    "    \n",
    "    plt.figure(num=1)\n",
    "    plt.figure(num=2)\n",
    "    plt.figure(num=3)\n",
    "    \n",
    "    \n",
    "    optimizer = torch.optim.Adam(net.parameters(), weight_decay=1e-4)\n",
    "#     optimizer = torch.optim.Adam(net.parameters(), lr=0.0005)\n",
    "        \n",
    "    \n",
    "    scheduler = torch.optim.lr_scheduler.ExponentialLR(optimizer, 0.95)\n",
    "    \n",
    "    \n",
    "    for epoch_index in tqdm(range(n_epochs)):\n",
    "        net.train()\n",
    "        print('Epoch: ' + str(epoch_index + 1) + ' of ' + str(N_EPOCHS))\n",
    "\n",
    "        tot_count = 0\n",
    "        tot_loss = 0\n",
    "        tot_accurate = 0\n",
    "        class_correct = list(0. for i in range(2))\n",
    "        class_total = list(0. for i in range(2))\n",
    "#         for batch_index, batch in enumerate(tqdm(data_loader)):\n",
    "        for batch in train_loader:\n",
    "            I1 = Variable(batch['I1'].float().cuda())\n",
    "            I2 = Variable(batch['I2'].float().cuda())\n",
    "            label = torch.squeeze(Variable(batch['label'].cuda()))\n",
    "\n",
    "            optimizer.zero_grad()\n",
    "            output = net(I1, I2)\n",
    "            loss = criterion(output, label.long())\n",
    "            loss.backward()\n",
    "            optimizer.step()\n",
    "            \n",
    "        scheduler.step()\n",
    "\n",
    "\n",
    "        epoch_train_loss[epoch_index], epoch_train_accuracy[epoch_index], cl_acc, pr_rec = test(train_dataset)\n",
    "        epoch_train_nochange_accuracy[epoch_index] = cl_acc[0]\n",
    "        epoch_train_change_accuracy[epoch_index] = cl_acc[1]\n",
    "        epoch_train_precision[epoch_index] = pr_rec[0]\n",
    "        epoch_train_recall[epoch_index] = pr_rec[1]\n",
    "        epoch_train_Fmeasure[epoch_index] = pr_rec[2]\n",
    "        \n",
    "#         epoch_test_loss[epoch_index], epoch_test_accuracy[epoch_index], cl_acc, pr_rec = test(test_dataset)\n",
    "        epoch_test_loss[epoch_index], epoch_test_accuracy[epoch_index], cl_acc, pr_rec = test(test_dataset)\n",
    "        epoch_test_nochange_accuracy[epoch_index] = cl_acc[0]\n",
    "        epoch_test_change_accuracy[epoch_index] = cl_acc[1]\n",
    "        epoch_test_precision[epoch_index] = pr_rec[0]\n",
    "        epoch_test_recall[epoch_index] = pr_rec[1]\n",
    "        epoch_test_Fmeasure[epoch_index] = pr_rec[2]\n",
    "\n",
    "        plt.figure(num=1)\n",
    "        plt.clf()\n",
    "        l1_1, = plt.plot(t[:epoch_index + 1], epoch_train_loss[:epoch_index + 1], label='Train loss')\n",
    "        l1_2, = plt.plot(t[:epoch_index + 1], epoch_test_loss[:epoch_index + 1], label='Test loss')\n",
    "        plt.legend(handles=[l1_1, l1_2])\n",
    "        plt.grid()\n",
    "#         plt.gcf().gca().set_ylim(bottom = 0)\n",
    "        plt.gcf().gca().set_xlim(left = 0)\n",
    "        plt.title('Loss')\n",
    "        display.clear_output(wait=True)\n",
    "        display.display(plt.gcf())\n",
    "\n",
    "        plt.figure(num=2)\n",
    "        plt.clf()\n",
    "        l2_1, = plt.plot(t[:epoch_index + 1], epoch_train_accuracy[:epoch_index + 1], label='Train accuracy')\n",
    "        l2_2, = plt.plot(t[:epoch_index + 1], epoch_test_accuracy[:epoch_index + 1], label='Test accuracy')\n",
    "        plt.legend(handles=[l2_1, l2_2])\n",
    "        plt.grid()\n",
    "        plt.gcf().gca().set_ylim(0, 100)\n",
    "#         plt.gcf().gca().set_ylim(bottom = 0)\n",
    "#         plt.gcf().gca().set_xlim(left = 0)\n",
    "        plt.title('Accuracy')\n",
    "        display.clear_output(wait=True)\n",
    "        display.display(plt.gcf())\n",
    "\n",
    "        plt.figure(num=3)\n",
    "        plt.clf()\n",
    "        l3_1, = plt.plot(t[:epoch_index + 1], epoch_train_nochange_accuracy[:epoch_index + 1], label='Train accuracy: no change')\n",
    "        l3_2, = plt.plot(t[:epoch_index + 1], epoch_train_change_accuracy[:epoch_index + 1], label='Train accuracy: change')\n",
    "        l3_3, = plt.plot(t[:epoch_index + 1], epoch_test_nochange_accuracy[:epoch_index + 1], label='Test accuracy: no change')\n",
    "        l3_4, = plt.plot(t[:epoch_index + 1], epoch_test_change_accuracy[:epoch_index + 1], label='Test accuracy: change')\n",
    "        plt.legend(handles=[l3_1, l3_2, l3_3, l3_4])\n",
    "        plt.grid()\n",
    "        plt.gcf().gca().set_ylim(0, 100)\n",
    "#         plt.gcf().gca().set_ylim(bottom = 0)\n",
    "#         plt.gcf().gca().set_xlim(left = 0)\n",
    "        plt.title('Accuracy per class')\n",
    "        display.clear_output(wait=True)\n",
    "        display.display(plt.gcf())\n",
    "\n",
    "        plt.figure(num=4)\n",
    "        plt.clf()\n",
    "        l4_1, = plt.plot(t[:epoch_index + 1], epoch_train_precision[:epoch_index + 1], label='Train precision')\n",
    "        l4_2, = plt.plot(t[:epoch_index + 1], epoch_train_recall[:epoch_index + 1], label='Train recall')\n",
    "        l4_3, = plt.plot(t[:epoch_index + 1], epoch_train_Fmeasure[:epoch_index + 1], label='Train Dice/F1')\n",
    "        l4_4, = plt.plot(t[:epoch_index + 1], epoch_test_precision[:epoch_index + 1], label='Test precision')\n",
    "        l4_5, = plt.plot(t[:epoch_index + 1], epoch_test_recall[:epoch_index + 1], label='Test recall')\n",
    "        l4_6, = plt.plot(t[:epoch_index + 1], epoch_test_Fmeasure[:epoch_index + 1], label='Test Dice/F1')\n",
    "        plt.legend(handles=[l4_1, l4_2, l4_3, l4_4, l4_5, l4_6])\n",
    "        plt.grid()\n",
    "        plt.gcf().gca().set_ylim(0, 1)\n",
    "#         plt.gcf().gca().set_ylim(bottom = 0)\n",
    "#         plt.gcf().gca().set_xlim(left = 0)\n",
    "        plt.title('Precision, Recall and F-measure')\n",
    "        display.clear_output(wait=True)\n",
    "        display.display(plt.gcf())\n",
    "        \n",
    "        \n",
    "#         mean_acc = (epoch_test_nochange_accuracy[epoch_index] + epoch_test_change_accuracy[epoch_index])/2\n",
    "#         if mean_acc > best_mean_acc:\n",
    "#             best_mean_acc = mean_acc\n",
    "#             save_str = 'net-best_epoch-' + str(epoch_index + 1) + '_acc-' + str(mean_acc) + '.pth.tar'\n",
    "#             torch.save(net.state_dict(), save_str)\n",
    "        \n",
    "        \n",
    "#         fm = pr_rec[2]\n",
    "        fm = epoch_train_Fmeasure[epoch_index]\n",
    "        if fm > best_fm:\n",
    "            best_fm = fm\n",
    "            save_str = 'net-best_epoch-' + str(epoch_index + 1) + '_fm-' + str(fm) + '.pth.tar'\n",
    "            torch.save(net.state_dict(), save_str)\n",
    "        \n",
    "        lss = epoch_train_loss[epoch_index]\n",
    "        if lss < best_lss:\n",
    "            best_lss = lss\n",
    "            save_str = 'net-best_epoch-' + str(epoch_index + 1) + '_loss-' + str(lss) + '.pth.tar'\n",
    "            torch.save(net.state_dict(), save_str)\n",
    "            \n",
    "            \n",
    "#         print('Epoch loss: ' + str(tot_loss/tot_count))\n",
    "        if save:\n",
    "            im_format = 'png'\n",
    "    #         im_format = 'eps'\n",
    "\n",
    "            plt.figure(num=1)\n",
    "            plt.savefig(net_name + '-01-loss.' + im_format)\n",
    "\n",
    "            plt.figure(num=2)\n",
    "            plt.savefig(net_name + '-02-accuracy.' + im_format)\n",
    "\n",
    "            plt.figure(num=3)\n",
    "            plt.savefig(net_name + '-03-accuracy-per-class.' + im_format)\n",
    "\n",
    "            plt.figure(num=4)\n",
    "            plt.savefig(net_name + '-04-prec-rec-fmeas.' + im_format)\n",
    "        \n",
    "    out = {'train_loss': epoch_train_loss[-1],\n",
    "           'train_accuracy': epoch_train_accuracy[-1],\n",
    "           'train_nochange_accuracy': epoch_train_nochange_accuracy[-1],\n",
    "           'train_change_accuracy': epoch_train_change_accuracy[-1],\n",
    "           'test_loss': epoch_test_loss[-1],\n",
    "           'test_accuracy': epoch_test_accuracy[-1],\n",
    "           'test_nochange_accuracy': epoch_test_nochange_accuracy[-1],\n",
    "           'test_change_accuracy': epoch_test_change_accuracy[-1]}\n",
    "    \n",
    "    print('pr_c, rec_c, f_meas, pr_nc, rec_nc')\n",
    "    print(pr_rec)\n",
    "    \n",
    "    return out\n",
    "\n",
    "L = 1024\n",
    "N = 2\n",
    "\n",
    "def test(dset):\n",
    "    net.eval()\n",
    "    tot_loss = 0\n",
    "    tot_count = 0\n",
    "    tot_accurate = 0\n",
    "    \n",
    "    n = 2\n",
    "    class_correct = list(0. for i in range(n))\n",
    "    class_total = list(0. for i in range(n))\n",
    "    class_accuracy = list(0. for i in range(n))\n",
    "    \n",
    "    tp = 0\n",
    "    tn = 0\n",
    "    fp = 0\n",
    "    fn = 0\n",
    "    \n",
    "    for img_index in dset.names:\n",
    "        I1_full, I2_full, cm_full = dset.get_img(img_index)\n",
    "        \n",
    "        s = cm_full.shape\n",
    "        \n",
    "\n",
    "        steps0 = np.arange(0,s[0],ceil(s[0]/N))\n",
    "        steps1 = np.arange(0,s[1],ceil(s[1]/N))\n",
    "        for ii in range(N):\n",
    "            for jj in range(N):\n",
    "                xmin = steps0[ii]\n",
    "                if ii == N-1:\n",
    "                    xmax = s[0]\n",
    "                else:\n",
    "                    xmax = steps0[ii+1]\n",
    "                ymin = jj\n",
    "                if jj == N-1:\n",
    "                    ymax = s[1]\n",
    "                else:\n",
    "                    ymax = steps1[jj+1]\n",
    "                I1 = I1_full[:, xmin:xmax, ymin:ymax]\n",
    "                I2 = I2_full[:, xmin:xmax, ymin:ymax]\n",
    "                cm = cm_full[xmin:xmax, ymin:ymax]\n",
    "\n",
    "                I1 = Variable(torch.unsqueeze(I1, 0).float()).cuda()\n",
    "                I2 = Variable(torch.unsqueeze(I2, 0).float()).cuda()\n",
    "                cm = Variable(torch.unsqueeze(torch.from_numpy(1.0*cm),0).float()).cuda()\n",
    "\n",
    "\n",
    "                output = net(I1, I2)\n",
    "                loss = criterion(output, cm.long())\n",
    "        #         print(loss)\n",
    "                tot_loss += loss.data * np.prod(cm.size())\n",
    "                tot_count += np.prod(cm.size())\n",
    "\n",
    "                _, predicted = torch.max(output.data, 1)\n",
    "\n",
    "                c = (predicted.int() == cm.data.int())\n",
    "                for i in range(c.size(1)):\n",
    "                    for j in range(c.size(2)):\n",
    "                        l = int(cm.data[0, i, j])\n",
    "                        class_correct[l] += c[0, i, j]\n",
    "                        class_total[l] += 1\n",
    "                        \n",
    "                pr = (predicted.int() > 0).cpu().numpy()\n",
    "                gt = (cm.data.int() > 0).cpu().numpy()\n",
    "                \n",
    "                tp += np.logical_and(pr, gt).sum()\n",
    "                tn += np.logical_and(np.logical_not(pr), np.logical_not(gt)).sum()\n",
    "                fp += np.logical_and(pr, np.logical_not(gt)).sum()\n",
    "                fn += np.logical_and(np.logical_not(pr), gt).sum()\n",
    "        \n",
    "    net_loss = tot_loss/tot_count\n",
    "    net_accuracy = 100 * (tp + tn)/tot_count\n",
    "    \n",
    "    for i in range(n):\n",
    "        class_accuracy[i] = 100 * class_correct[i] / max(class_total[i],0.00001)\n",
    "\n",
    "    prec = tp / (tp + fp)\n",
    "    rec = tp / (tp + fn)\n",
    "    f_meas = 2 * prec * rec / (prec + rec)\n",
    "    prec_nc = tn / (tn + fn)\n",
    "    rec_nc = tn / (tn + fp)\n",
    "    \n",
    "    pr_rec = [prec, rec, f_meas, prec_nc, rec_nc]\n",
    "        \n",
    "    return net_loss, net_accuracy, class_accuracy, pr_rec\n",
    "    \n",
    "    \n",
    "\n",
    "\n",
    "\n",
    "    \n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "if LOAD_TRAINED:\n",
    "    net.load_state_dict(torch.load('net_final.pth.tar'))\n",
    "    print('LOAD OK')\n",
    "else:\n",
    "    t_start = time.time()\n",
    "    out_dic = train()\n",
    "    t_end = time.time()\n",
    "    print(out_dic)\n",
    "    print('Elapsed time:')\n",
    "    print(t_end - t_start)\n",
    "\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "if not LOAD_TRAINED:\n",
    "    torch.save(net.state_dict(), 'net_final.pth.tar')\n",
    "    print('SAVE OK')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "\n",
    "\n",
    "def save_test_results(dset):\n",
    "    for name in tqdm(dset.names):\n",
    "        with warnings.catch_warnings():\n",
    "            I1, I2, cm = dset.get_img(name)\n",
    "            I1 = Variable(torch.unsqueeze(I1, 0).float()).cuda()\n",
    "            I2 = Variable(torch.unsqueeze(I2, 0).float()).cuda()\n",
    "            out = net(I1, I2)\n",
    "            _, predicted = torch.max(out.data, 1)\n",
    "            I = np.stack((255*cm,255*np.squeeze(predicted.cpu().numpy()),255*cm),2)\n",
    "            io.imsave(f'{net_name}-{name}.png',I)\n",
    "\n",
    "\n",
    "\n",
    "t_start = time.time()\n",
    "# save_test_results(train_dataset)\n",
    "save_test_results(test_dataset)\n",
    "t_end = time.time()\n",
    "print('Elapsed time: {}'.format(t_end - t_start))\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "L = 1024\n",
    "\n",
    "def kappa(tp, tn, fp, fn):\n",
    "    N = tp + tn + fp + fn\n",
    "    p0 = (tp + tn) / N\n",
    "    pe = ((tp+fp)*(tp+fn) + (tn+fp)*(tn+fn)) / (N * N)\n",
    "    \n",
    "    return (p0 - pe) / (1 - pe)\n",
    "\n",
    "def test(dset):\n",
    "    net.eval()\n",
    "    tot_loss = 0\n",
    "    tot_count = 0\n",
    "    tot_accurate = 0\n",
    "    \n",
    "    n = 2\n",
    "    class_correct = list(0. for i in range(n))\n",
    "    class_total = list(0. for i in range(n))\n",
    "    class_accuracy = list(0. for i in range(n))\n",
    "    \n",
    "    tp = 0\n",
    "    tn = 0\n",
    "    fp = 0\n",
    "    fn = 0\n",
    "    \n",
    "    for img_index in tqdm(dset.names):\n",
    "        I1_full, I2_full, cm_full = dset.get_img(img_index)\n",
    "        \n",
    "        s = cm_full.shape\n",
    "        \n",
    "        for ii in range(ceil(s[0]/L)):\n",
    "            for jj in range(ceil(s[1]/L)):\n",
    "                xmin = L*ii\n",
    "                xmax = min(L*(ii+1),s[1])\n",
    "                ymin = L*jj\n",
    "                ymax = min(L*(jj+1),s[1])\n",
    "                I1 = I1_full[:, xmin:xmax, ymin:ymax]\n",
    "                I2 = I2_full[:, xmin:xmax, ymin:ymax]\n",
    "                cm = cm_full[xmin:xmax, ymin:ymax]\n",
    "\n",
    "                I1 = Variable(torch.unsqueeze(I1, 0).float()).cuda()\n",
    "                I2 = Variable(torch.unsqueeze(I2, 0).float()).cuda()\n",
    "                cm = Variable(torch.unsqueeze(torch.from_numpy(1.0*cm),0).float()).cuda()\n",
    "\n",
    "                output = net(I1, I2)\n",
    "                    \n",
    "                loss = criterion(output, cm.long())\n",
    "                tot_loss += loss.data * np.prod(cm.size())\n",
    "                tot_count += np.prod(cm.size())\n",
    "\n",
    "                _, predicted = torch.max(output.data, 1)\n",
    "\n",
    "                c = (predicted.int() == cm.data.int())\n",
    "                for i in range(c.size(1)):\n",
    "                    for j in range(c.size(2)):\n",
    "                        l = int(cm.data[0, i, j])\n",
    "                        class_correct[l] += c[0, i, j]\n",
    "                        class_total[l] += 1\n",
    "                        \n",
    "                pr = (predicted.int() > 0).cpu().numpy()\n",
    "                gt = (cm.data.int() > 0).cpu().numpy()\n",
    "                \n",
    "                tp += np.logical_and(pr, gt).sum()\n",
    "                tn += np.logical_and(np.logical_not(pr), np.logical_not(gt)).sum()\n",
    "                fp += np.logical_and(pr, np.logical_not(gt)).sum()\n",
    "                fn += np.logical_and(np.logical_not(pr), gt).sum()\n",
    "        \n",
    "    net_loss = tot_loss/tot_count        \n",
    "    net_loss = float(net_loss.cpu().numpy())\n",
    "    \n",
    "    net_accuracy = 100 * (tp + tn)/tot_count\n",
    "    \n",
    "    for i in range(n):\n",
    "        class_accuracy[i] = 100 * class_correct[i] / max(class_total[i],0.00001)\n",
    "        class_accuracy[i] =  float(class_accuracy[i].cpu().numpy())\n",
    "\n",
    "    prec = tp / (tp + fp)\n",
    "    rec = tp / (tp + fn)\n",
    "    dice = 2 * prec * rec / (prec + rec)\n",
    "    prec_nc = tn / (tn + fn)\n",
    "    rec_nc = tn / (tn + fp)\n",
    "    \n",
    "    pr_rec = [prec, rec, dice, prec_nc, rec_nc]\n",
    "    \n",
    "    k = kappa(tp, tn, fp, fn)\n",
    "    \n",
    "    return {'net_loss': net_loss, \n",
    "            'net_accuracy': net_accuracy, \n",
    "            'class_accuracy': class_accuracy, \n",
    "            'precision': prec, \n",
    "            'recall': rec, \n",
    "            'dice': dice, \n",
    "            'kappa': k}\n",
    "\n",
    "results = test(test_dataset)\n",
    "pprint(results)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.7.6"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}
